#!/usr/bin/python

'''

Description: 
    This script is to filter out less occurrent 
    topics from the beta matrix (an output of LDA) 

Created on: 
    Jun 27, 2011

Author: 
    Clint P. George 

'''

import os
import glob
import subprocess

from numpy import *


if __name__ == '__main__':
    
    batch_dir = raw_input('Batch folder: ') # '/home/clint/Dropbox/TREC/batch'
    saved_beta = raw_input('Beta file extension: ')  # 'hdp-topics.dat'
    
    
    for root, dirs, files in os.walk(batch_dir):
        
        print 'reading ', root, '...' 
        
        for file in files:            
            if file == saved_beta:

                with open(os.path.join(root, file), 'r') as fbeta:
                    ml = []
                    for each_line in fbeta:
                        ll = each_line.strip().split(' ')                        
                        ll2 = [int(x) for x in ll]
                        ml.append(ll2)

                m = array(ml)
                row_sums = m.sum(axis=1)
                row_mean = row_sums.mean() 
                
                vaild_topics = ones(len(row_sums), dtype=int16)  
                idx = 0
                for s in row_sums: 
                    if s < (row_mean * 0.1):
                        vaild_topics[idx] = 0
                    idx += 1

                with open(os.path.join(root, file), 'r') as fbeta:
                    with open(os.path.join(root, 'beta_counts.dat'), 'w') as ff:
                        with open(os.path.join(root, 'vaild_topic_idx'), 'w') as fv:
                            idx = 0 
                            for each_line in fbeta:
                                fv.write(str(vaild_topics[idx]) + '\n')
                                if vaild_topics[idx] == 1:
                                    ff.write(each_line.strip() + '\n')
                                idx += 1 
